#!/usr/bin/env python3
"""
YOLO模型训练脚本
使用ultralytics库训练YOLO模型
"""

import os
import yaml
import argparse
from pathlib import Path
from ultralytics import YOLO

def train_yolo_model(dataset_config, model_size="n", epochs=100, batch_size=16, img_size=640):
    """训练YOLO模型"""
    
    # 创建YOLO模型
    model = YOLO(f"yolo{model_size}.pt")  # 加载预训练模型
    
    # 训练参数
    train_args = {
        'data': dataset_config,  # 数据集配置文件路径
        'epochs': epochs,        # 训练轮数
        'batch': batch_size,     # 批次大小
        'imgsz': img_size,       # 图像尺寸
        'patience': 20,          # 早停耐心值
        'save': True,            # 保存模型
        'save_period': 10,       # 每10个epoch保存一次
        'cache': False,          # 是否缓存图像
        'device': 'auto',        # 自动选择设备
        'workers': 8,            # 数据加载器工作进程数
        'project': 'runs/train', # 项目名称
        'name': 'plant_diagnosis', # 实验名称
        'exist_ok': True,        # 允许覆盖现有实验
        'pretrained': True,      # 使用预训练权重
        'optimizer': 'auto',     # 优化器
        'verbose': True,         # 详细输出
        'seed': 42,              # 随机种子
        'deterministic': True,   # 确定性训练
        'single_cls': False,     # 多类别
        'rect': False,           # 矩形训练
        'cos_lr': False,         # 余弦学习率调度
        'close_mosaic': 10,      # 最后10个epoch关闭mosaic
        'resume': False,         # 是否恢复训练
        'amp': True,             # 混合精度训练
        'fraction': 1.0,         # 数据集比例
        'profile': False,        # 性能分析
        'freeze': None,          # 冻结层
        'lr0': 0.01,            # 初始学习率
        'lrf': 0.01,            # 最终学习率
        'momentum': 0.937,       # SGD动量
        'weight_decay': 0.0005,  # 权重衰减
        'warmup_epochs': 3.0,    # 预热轮数
        'warmup_momentum': 0.8,  # 预热动量
        'warmup_bias_lr': 0.1,   # 预热偏置学习率
        'box': 7.5,              # 框损失增益
        'cls': 0.5,              # 分类损失增益
        'dfl': 1.5,              # DFL损失增益
        'pose': 12.0,            # 姿态损失增益
        'kobj': 2.0,             # 关键点目标损失增益
        'label_smoothing': 0.0,  # 标签平滑
        'nbs': 64,               # 标称批次大小
        'overlap_mask': True,    # 掩码重叠
        'mask_ratio': 4,         # 掩码下采样比率
        'dropout': 0.0,          # 使用分类器dropout
        'val': True,             # 验证
        'plots': True,           # 保存图表
    }
    
    # 开始训练
    print(f"开始训练YOLO模型...")
    print(f"数据集配置: {dataset_config}")
    print(f"模型大小: YOLO{model_size.upper()}")
    print(f"训练轮数: {epochs}")
    print(f"批次大小: {batch_size}")
    print(f"图像尺寸: {img_size}")
    
    results = model.train(**train_args)
    
    print(f"训练完成！")
    print(f"最佳模型保存在: {results.save_dir}")
    
    return results

def validate_model(model_path, dataset_config):
    """验证模型性能"""
    model = YOLO(model_path)
    
    # 在验证集上评估
    results = model.val(data=dataset_config)
    
    print(f"验证结果:")
    print(f"mAP50: {results.box.map50:.4f}")
    print(f"mAP50-95: {results.box.map:.4f}")
    print(f"Precision: {results.box.mp:.4f}")
    print(f"Recall: {results.box.mr:.4f}")
    
    return results

def main():
    parser = argparse.ArgumentParser(description="训练YOLO模型")
    parser.add_argument("--dataset", type=str, required=True, help="数据集配置文件路径")
    parser.add_argument("--model-size", type=str, default="n", choices=["n", "s", "m", "l", "x"], help="模型大小")
    parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
    parser.add_argument("--batch-size", type=int, default=16, help="批次大小")
    parser.add_argument("--img-size", type=int, default=640, help="图像尺寸")
    parser.add_argument("--validate", action="store_true", help="训练后验证模型")
    
    args = parser.parse_args()
    
    # 检查数据集配置文件
    if not os.path.exists(args.dataset):
        print(f"错误: 数据集配置文件不存在: {args.dataset}")
        return
    
    # 训练模型
    results = train_yolo_model(
        dataset_config=args.dataset,
        model_size=args.model_size,
        epochs=args.epochs,
        batch_size=args.batch_size,
        img_size=args.img_size
    )
    
    # 验证模型
    if args.validate:
        best_model_path = results.save_dir / "weights" / "best.pt"
        if best_model_path.exists():
            print(f"\n验证最佳模型: {best_model_path}")
            validate_model(best_model_path, args.dataset)
        else:
            print(f"警告: 最佳模型文件不存在: {best_model_path}")

if __name__ == "__main__":
    main()
